import pickle
import multiprocessing as mp
import trimesh
from autolab_core import RigidTransform
from copy import deepcopy
import numpy as np 
import json
import os
#from .grasping.grasp import ParallelJawPtGrasp3D
from grasping.grasp_sampler import GraspSampler

#def _get_rotation_matrix_y(theta):
#    cos_t = np.cos(theta)
#    sin_t = np.sin(theta)
#    R = nnp.c_[[cos_t, 0, sin_t], np.c_[[0, 1, 0], [-sin_t, 0, cos_t]]]
#    return R
#def unrotated_full_axis(axis):
#    grasp_axis_y = axis
#    grasp_axis_x = np.array([grasp_axis_y[1], -grasp_axis_y[0], 0])
#    if np.linalg.norm(grasp_axis_x) == 0:
#        grasp_axis_x = np.array([1,0,0])
#    grasp_axis_x = grasp_axis_x / norm(grasp_axis_x)
#    grasp_axis_z = np.cross(grasp_axis_x, grasp_axis_y)
#    R = np.c_[grasp_axis_x, np.c_[grasp_axis_y, grasp_axis_z]]
#    return R
#
#def rotated_full_axis(approach_angle, axis)
#    R = _get_rotation_matrix_y(approach_angle)
#    R = unrotated_full_axis(axis).dot(R)
#    return R
def collision_checker(scene, mesh, finger_l, grasp, R_l, R_r):
#    scene = trimesh.Scene([mesh])
    
    collision_manager, _ = trimesh.collision.scene_to_collision(scene)
    scene_ = deepcopy(scene)
    finger_left = deepcopy(finger_l)
    finger_right = deepcopy(finger_l)
    T_finger = trimesh.transformations.translation_matrix([-(grasp.real_jaw_width_/2+0.005), 0,0])
    finger_left.apply_transform(T_finger)
    finger_right.apply_transform(T_finger)
    finger_left.apply_transform(R_l)
    finger_right.apply_transform(R_r)
    T_leftfinger = RigidTransform(grasp.rotated_full_axis, grasp.contactpoints[0], from_frame='left_finger', to_frame='obj')
    T_rightfinger = RigidTransform(grasp.rotated_full_axis, grasp.contactpoints[1], from_frame='right_finger', to_frame='obj')
    finger_left.apply_transform(T_leftfinger.matrix)
    finger_right.apply_transform(T_rightfinger.matrix)
    is_collision = collision_manager.in_collision_single(finger_left, return_names=False) or collision_manager.in_collision_single(finger_right, return_names=False)

    if debug_vis:
        scene_.add_geometry(finger_left)
        scene_.add_geometry(finger_right)
        scene_.show()
    #if is_collision:
    #    if debug_vis:
    #        scene_.add_geometry(finger_left)
    #        scene_.add_geometry(finger_right)
    #        scene_.show()

    return is_collision

def prune_grasp_method(dataset_dict, R_l, R_r, key):
        path = dataset_dict[key]['path']
        mesh = trimesh.load_mesh(path)
        files = os.listdir('/home/v-wewei/code/two_stage_pointnet/generated_grasp_2500/{}'.format(key))

        for file_name in files:

            if os.path.splitext(file_name)[0].startswith('pruned'):
                with open('/home/v-wewei/code/two_stage_pointnet/generated_grasp_2500/{}/{}'.format(key, file_name), 'rb') as f:
                    grasps = pickle.load(f)
                            

                    selected_grasps = []
                    scene = trimesh.Scene([mesh])
                    for grasp in grasps:
                        is_collision = collision_checker(scene, mesh, finger_l, grasp, R_l, R_r)
                        if not is_collision:
                            selected_grasps.append(grasp)
                    with open('/home/v-wewei/code/two_stage_pointnet/generated_grasp_2500/{}/Pruned_{}.pickle'.format(key, len(selected_grasps)), 'wb') as f_1:
                        pickle.dump(selected_grasps, f_1)
        print('finish')

if __name__ == "__main__":
    debug_vis = False
    T = trimesh.transformations.translation_matrix([-0.0, -0.0065, -0.048])
    R_l = trimesh.transformations.euler_matrix(np.pi/2, np.pi/2, 0, 'rxyz')
    R_r = trimesh.transformations.euler_matrix(np.pi/2, np.pi/2, np.pi, 'rxyz')
    base_path = '/home/v-wewei/code/two_stage_pointnet/data/grippers/yumi/coarse/base.stl'
    base = trimesh.load_mesh(base_path)
    T_base = trimesh.transformations.translation_matrix([0, 0, -0.13])
    R_base = trimesh.transformations.euler_matrix(np.pi/2, np.pi/2, 0,'rxyz')
    base.apply_transform(T_base)
    base.apply_transform(R_base)
    finger_l = trimesh.load('/home/v-wewei/code/two_stage_pointnet/data/grippers/yumi/coarse/finger.stl')
    finger_l.apply_transform(T)
    pool = mp.Pool(processes=mp.cpu_count())
    with open('/home/v-wewei/grasp_sampler/dataset_dict.json', 'r') as f:
        dataset_dict = json.load(fp=f)
        for key in dataset_dict.keys():
#            if key != '024_bowl':
#                continue
            pool.apply_async(prune_grasp_method, args=(dataset_dict, R_l, R_r, key,))
            #prune_grasp_method(dataset_dict, R_l, R_r, key)

        pool.close()
        pool.join()
